const fs = require("fs");
const tf = require("@tensorflow/tfjs-node");
const img2x = (buffer) => {
    return tf.tidy(() => {
        //buffer转化为tensor
        const imgTs = tf.node.decodeImage(new Uint8Array(buffer));
        //改变图片的大小统一尺寸
        const imgTsResized = tf.image.resizeBilinear(imgTs, [224, 224]);
        //将图片归一化，在指定通道
        return imgTsResized.toFloat().sub(255 / 2).div(255 / 2).reshape([1, 224, 224, 3]);
    })
}
const getData = async (trainDir, outputDir) => {
    console.log("getData");
    const classes = fs.readdirSync(trainDir).filter(n => !n.includes("."));
    fs.writeFileSync(`${outputDir}/classes.json`, JSON.stringify(classes));

    const inputs = [];
    const lables = [];
    classes.forEach((dir, dirIndex) => {
        fs.readdirSync(`${trainDir}/${dir}`)
            .filter((n) => n.match(/jpg$/))
            .slice(0,100)
            .forEach((filename) => {
                console.log("读取:", dir, filename);
                const imgPath = `${trainDir}/${dir}/${filename}`;
                const buffer = fs.readFileSync(imgPath);
                const x = img2x(buffer);
                inputs.push(x);
                lables.push(dirIndex);
            })
    })

    const xs = tf.concat(inputs);
    const ys = tf.tensor(lables);
    return {
        xs,ys,classes
    };
}

module.exports = getData;